#ifndef __HIRCODE_H__
#define __HIRCODE_H__

#include "HVisitor.h"
//#include "HSymbolTable.h"
#include "HFunctionTable.h"
#include "RV64.h"
#include <iostream>
#include <map>
#include <set>

// base types, or `VOID` for label
enum class HIRBtype {
	/*
	 * Make sure that the values of the four base types
	 * and `VOID`, `UNKNOWN` are the same as those in `Btype`!
	 */
	UNKNOWN = int(Btype::UNKNOWN),
	VOID = int(Btype::VOID),
	I32 = int(Btype::INT),
	I8 = int(Btype::BOOL),
	FLOAT = int(Btype::FLOAT),
	DOUBLE = int(Btype::DOUBLE),
	POINTER,
	MAX
};

// Convert `Btype` to `HIRBtype`, and `HIRBtype::UNKNOWN` if failed
static inline HIRBtype CACTBtype2HIRBtype(Btype btype)
{
	if (unsigned(btype) <= unsigned(Btype::DOUBLE))
		return HIRBtype(unsigned(btype));
	else
		return HIRBtype::UNKNOWN;
}

// type for HXD1 IR temporary number (index)
typedef uint32_t HIRTempno;
#define INVALID_TEMP_NO UINT32_MAX
// type for HXD1 IR labels
typedef HIRTempno HIRLabel;
#define INVALID_LABEL INVALID_TEMP_NO
// type for HXD1 IR immediate
typedef union {
	int64_t ival;
	float fval;
	double dval;
} HIRImm;

/*
 * For HXD1 IMM `imm` with base type `btype`, get its binary
 * representation (sign-extened to 64-bit).
 * Usually for appending RV64 `li` instruction.
 */
static inline int64_t hirimm2int64(const HIRImm &imm, HIRBtype btype)
{
	const float *pfval = &imm.fval;
	int64_t ret = 0l;
	if (btype == HIRBtype::DOUBLE)
		ret = imm.ival;
	else if (btype == HIRBtype::I32)
		ret = int32_t(imm.ival);
	else if (btype == HIRBtype::I8)
		ret = int8_t(imm.ival);
	else if (btype == HIRBtype::FLOAT)
		ret = *(const int32_t*)(pfval);
	return ret;
}

// to set a local variable that is not initialized
#define IMM_GARBAGE 0x55555555

#define INVALID_BB_IDX UINT32_MAX

// HXD1 IR Temporary
struct HIRTemp {
	/*
	 * Type of this temporary.
	 * If it is a label, `type` is `HIRBtype::VOID`.
	 */
	HIRBtype type;
	/*
	 * Indexes of the first instruction that assigns this temporary,
	 * the first and the last instruction that refereces it.
	 * It is for future analysis and optimization, such as
	 * registers allocation.
	 * NOTE: if the instructions are reordered for optimization,
	 * remember to update this field.
	 */
	uint32_t fassigned, frefed, lrefed;
#define INVALID_INST_IDX UINT32_MAX
	// Is it a local temp of a BB? `INVALID_BB_IDX` is for non-local temp.
	uint32_t localof;
	// Live flag for dead code elimination
	bool live;
	// Spill cost
	double spillcost;
	// A constructor to ban default constructing
	HIRTemp(HIRBtype tp)
	{
		type = tp;
		fassigned = frefed = lrefed = INVALID_INST_IDX;
		localof = INVALID_BB_IDX; live = false;
		spillcost = 0.0;
	}
};

// moved from class `HIRCode`
struct extvar_t {
	// symbol name
	std::string id;
	// in `.rodata` section or `.data`/`.bss`?
	bool is_const;
	// is array?
	bool is_array;
	// base type
	HIRBtype btype;
	// number of element
	size_t length;
	/*
	 * initialzation data:
	 * Hold the initialization data of the external variable, or
	 * `nullptr` if it should be filled with zero.
	 * `const` may be deleted in the future if needed.
	 */
	const void *data;
	/*
	 * Alignment (byte). Temporarily, set alignment for
	 * arraies as 16 B, and other variables as the max of 4 B and its size.
	 */
	size_t alignment;
};

/*
 * Type of data in HXD1D IR.
 * For example:
 * `{1u, HIRBtype::I32}` represents `i32`;
 * `{5u, HIRBtype::FLOAT}` represents `[5 x float]`.
 */
struct HIRType {
	// for array
	uint32_t length;
	/*
	 * `Btype` is corresponding with HXD1D IR type:
	 * `BOOL`: bool;
	 * `INT`: i32;
	 * `FLOAT`: float;
	 * `DOUBLE`: double;
	 */
	HIRBtype base_type;
};
/*
 * Get the size of a base type.
 * NOTE: size of a base type must be power of 2,
 * which is used in HXD1 back end (RV64 generator).
 */
static inline uint32_t hirsize(const HIRBtype &btype)
{
	uint32_t sz;
	switch (btype)
	{
	case HIRBtype::I32: case HIRBtype::FLOAT:
		sz = 4u; break;
	case HIRBtype::I8:
		sz = 1u; break;
	case HIRBtype::DOUBLE: case HIRBtype::POINTER:
		sz = 8u; break;
	default:
		sz = 0u; break;
	}
#define HIRBTYPE_SIZE_MAX 8u
	return sz;
}
// Get the size of a type (multiplied by its length)
static inline size_t hirsize(const HIRType &type)
{
	return hirsize(type.base_type) * size_t(type.length);
}

/*
 * Operand for HXD1 IR, that can be
 * a temporary, an immediate, or
 * a pointer to external variable in IR table.
 */
struct HIROpnd {
	enum {TEMP, IMM, XVAR} tp;
	union {	// 8 bytes
		HIRTempno temp;
		HIRImm imm;
		// external variable index
		size_t xvaridx;
	};
	bool operator==(const HIROpnd &op) const
	{
		if (tp != op.tp) return false;
		if (tp == HIROpnd::TEMP)
			return temp == op.temp;
		else if (tp == HIROpnd::IMM)
		/*
		 * NOTE: type of the immediate is ignored!
		 * So if two `i32` imm is the same in the lower 32 bits
		 * but different in the higher 32 bits, they will be considered not equal.
		 */
			return imm.ival == op.imm.ival;
		else if (tp == HIROpnd::XVAR)
			return xvaridx == op.xvaridx;
		else
		// Wrong type!
			return false;
	}	
};

// zero immediate
const HIROpnd Zero{HIROpnd::IMM, 0ll};

/*
 * Function argument for HXD1 IR.
 */
struct HIRFarg {
	HIRBtype type;
	HIROpnd arg;
};

// IR operations
enum class HIROper {
	RET, BRc, BRu,	// conditional/un-
	FNEG,
	ADD, FADD, SUB, FSUB, MUL, FMUL, UDIV, SDIV, FDIV, UREM, SREM,
	SHL, LSHR, ASHR, AND, OR, XOR,
	ALLOCA, LOAD, STORE,
	ICMP, FCMP,
	PHI, CALL,
	LABEL,
	EMPTY_LINE		// for an empty line
};

// `cond` for `icmp` instruction
enum class ICMPCond {
	EQ, NE, SGT, SGE, SLT, SLE
};

// `cond` for `fcmp` instruction
enum class FCMPCond {
	OEQ, ONE, OGT, OGE, OLT, OLE
};

struct BasicBlock {
	uint32_t begin, end;
	std::vector<uint32_t> pred, succ;
	// Should be initialized as `false` and set while traversing
	bool reachable;
	/*
	 * Temporaries that are assigned in this BB.
	 * `asstemps[#t]` is the index of the first instruction in this block
	 * that assigns temp `#t`.
	 */
	std::map<HIRTempno, uint32_t> asstemps;
	// Sets to solve flow equations
	std::set<HIRTempno> livevar_def, livevar_use;
	std::set<HIRTempno> livevar_in, livevar_out;
	// Dominator BB
	std::set<uint32_t> dominators;
// moved to the beginning
//#define INVALID_BB_IDX UINT32_MAX
};

// Forward declaration
struct HIRInst;

class HIRFunct {
public:
	// function name
	std::string id;
	// return value type
	HIRBtype rettype;
	/*
	 * Number of parameters of the function.
	 * NOTE: the first `nparams` `temps` are bound with arguments.
	 */
	uint32_t nparams;
	/*
	 * Temporaries used in this function,
	 * whose index is the temporary number.
	 */
	std::vector<HIRTemp> temps;
	// number of `ALLOCA` instructions at the beginning
	uint32_t nalloca;
	// HXD1 IR instructions
	std::vector<HIRInst> ircode;
	// Basic blocks
	std::vector<BasicBlock> basic_blocks;
	// Loops
	std::vector<std::set<uint32_t>> loops;
	// constructor
	HIRFunct()
	{
		this->rettype = HIRBtype::UNKNOWN;
		this->nalloca = 0u;
		this->nparams = 0u;
	}
	size_t insert_inst(const HIRInst &inst, size_t loc);
	int opt_remove_redundant_labels(void);
	void opt_control_flow_graph_cleanup(void);
	uint32_t divide_basic_blocks(void);
	void construct_control_flow_graph(void);
	void traverse_control_flow_graph(BasicBlock & bb);
	void opt_strength_reduction(void);
	void opt_merge_adjacent(void);
	size_t opt_dead_code_elimination(bool use_live_temps = false);
	size_t opt_common_subexpr_elimination(bool use_live_temps = false);
	void dfa_solve_livevar(void);
	void dfa_solve_dominators(void);
	void dfa_detect_loops(void);
	void dfa_calc_spillcost(void);

	// Functions below are used to support code generator

	uint64_t get_temp_future_use(HIRTempno t, uint32_t start_iidx, uint32_t bbidx) const;
};

// HXD1 IR Instruction
struct HIRInst {
	union {
		// Destination that can only be a temporary
		HIRTempno dest;
		// Label number for `LABEL`
		HIRLabel labelno;
	};
	// Operation
	HIROper oper;
	// Source type, only scalar except for `ALLOCA`
	HIRType type;
	// Source 0
	union {
		HIRTempno cond;			// `BRc`
		HIROpnd op1;			// `fneg`, `add`, ..., `xor`, `icmp`, `fcmp`
		HIROpnd pointer;		// `load`, `store`
		HIROpnd val0;			// `phi`
		size_t funcidx;			// `call` // function index in `_funct` of `HIRCode`
	};
	// Source 1
	union {
		HIROpnd value;		// `ret`, `store`
		HIRLabel iftrue;	// `BRc`
		HIRLabel brdest;	// `BRu`
		HIROpnd op2;		// `add`, ..., `xor`, `icmp`, `fcmp`
		HIRLabel label0;	// `phi`
	};
	// Source 2
	union {
		HIRLabel iffalse;	// `BRc`
		uint32_t alignment;	// for `alloca`, only power of 2
		ICMPCond icond;		// `icmp`
		FCMPCond fcond;		// `fcmp`
		HIROpnd val1;		// `phi`
		/*
		 * Flag of garbage initialization.
		 * Only for `add` and `fadd`!
		 */
		bool flag_garbage;
	};
	// Source 3
	union {
		HIRLabel label1;				// `phi`
		std::vector<HIRFarg> *pfargs;	// `call`, may be `nullptr` if no arg
	};
	/*
	 * Some info can be attached to an instruction
	 * that will be dumped in IR code file for debug.
	 * NOTE: it must be created using `new`.
	 */
	std::string *cominfo;
	HIRInst(HIROper oper_ = HIROper::LABEL,
		HIRTempno dest_ = INVALID_TEMP_NO) {
		dest = dest_;
		oper = oper_;
		flag_garbage = false;
		type.length = 1u;
		pfargs = nullptr;
		cominfo = nullptr;
	}
	HIRTempno get_dest(void) const;
	size_t get_srctemps(std::vector<HIRTempno> &srctemps) const;
	bool opequals(const HIRInst &inst, bool ignore_side_effect = false) const;
	size_t srctemps_replace(HIRTempno rip, HIRTempno rep);
	void brdests_replace(HIRLabel rip, HIRLabel rep);
	HIRTempno is_copy(void) const;
};

class HIRCode {
	static unsigned int ymr;
private:
	// external variables
	std::vector<extvar_t> _extvar;
	// functions defined by user
	std::vector<HIRFunct> _funct;
	// the first `_nbuiltin` functions in `_funct` are builtin
	size_t _nbuiltin;
	/*
	 * Pointer to current function determining to which function
	 * IR instructions will be added.
	 */
	size_t _current_function;
	// ...
public:
	HIRCode();
	~HIRCode();

	void add_builtin(const HFunctionTable &ftable, size_t nbuiltin);
	
	void exit_function() { _current_function = SIZE_MAX; }
	size_t enter_function(size_t funcidx);

	size_t nbuiltin() const { return _nbuiltin; }
	const std::vector<extvar_t>& extvars() const { return _extvar; }
	const extvar_t & extvar(size_t xvaridx) const { return _extvar.at(xvaridx); }

	const std::vector<HIRFunct>& functs() const { return _funct; }
	const HIRFunct & funct(size_t functidx) const { return _funct.at(functidx); }

	size_t add_extvar(const HSTEntry &hstentry, const void *data);
	size_t add_funct(const HFTEntry & hftentry,
		HSymbolTable *psymtb, const std::vector<size_t>& params);
	size_t search_extvar(const std::string &id) const;
	size_t search_funct(const std::string &id) const;

	const HIRTemp & temp(HIRTempno tempno) const
		{ return _funct.at(_current_function).temps.at(tempno); }
	const std::vector<HIRTemp> & temps() const
		{ return _funct.at(_current_function).temps; }
	// function to create temporary
	HIRTempno new_temp(HIRBtype type);
	void update_temp_fassigned(HIRTempno tempno, uint32_t aiidx);
	void update_temp_flrefed(HIRTempno tempno, uint32_t riidx);

	//uint64_t generate_init_code(HSymbolTable *phsymtb, size_t hsteidx, HIRTempno ti);
	uint64_t generate_init_code(HSymbolTable *phsymtb, size_t hsteidx,
		const void *init_data);
	HIROpnd generate_expr_code(const HExpression *expr, bool as_cond = false,
		HIRLabel lbtrue = INVALID_LABEL, HIRLabel lbfalse = INVALID_LABEL,
		bool as_lval = false);
	void post_process(void);
	
	// functions to add instruction
	size_t append_inst(const HIRInst &inst);
	size_t insert_inst(const HIRInst &inst, size_t loc);
	/*
	 * Append an instruction and return its index.
	 * `lassigned` of `dest` may be set automatically.
	 */
	size_t append_ret(HIRBtype type, HIROpnd value);
	size_t append_br(HIRTempno cond, HIRLabel iftrue, HIRLabel iffalse);
	size_t append_br(HIRLabel dest);
	size_t append_fneg(HIRTempno dest, HIRBtype type, HIROpnd op1);
	size_t append_arlo(HIROper op, HIRTempno dest, HIRBtype type,
		const HIROpnd &op1, const HIROpnd &op2, const char *cominfo = nullptr,
		bool flag_garbage = false);
	size_t insert_alloca(HIRTempno dest, const HIRType &type, uint32_t alignment,
		const char *cominfo = nullptr);
	size_t append_load(HIRTempno dest, HIRBtype type, HIRTempno ptr);
	size_t append_load(HIRTempno dest, HIRBtype type, size_t xvidx);
	size_t append_store(HIRBtype type, const HIROpnd &ptr, const HIROpnd &value,
		const char *cominfo = nullptr);
	size_t append_icmp(HIRTempno dest, HIRBtype type, ICMPCond cond,
		const HIROpnd &op1, const HIROpnd &op2);
	size_t append_fcmp(HIRTempno dest, HIRBtype type, FCMPCond cond,
		const HIROpnd &op1, const HIROpnd &op2);
	size_t append_phi(HIRTempno dest, HIRBtype type,
		HIRLabel label0, const HIROpnd &val0, HIRLabel label1, const HIROpnd &val1);
	size_t append_call(HIRBtype type, size_t funcidx, std::vector<HIRFarg> *pfargs,
		HIRTempno dest = INVALID_TEMP_NO);
	size_t append_label(HIRLabel labelno);
	size_t append_emptylines(uint32_t nel = 2u);

	void dump(std::ostream &os, const char *src_fname = nullptr) const;
};

extern HIRCode hxd1_ir;

/*
 * Copy a base type scalar value from `src` to `dst`.
 * If type `ty` is invalid, 1 will be returned.
 */
static inline int elecpy(void *dst, const void *src, Btype ty)
{
	switch (ty)
	{
	case Btype::FLOAT: case Btype::INT:
		*((int32_t*)(dst)) = *((int32_t*)(src));
		break;
	case Btype::DOUBLE:
		*((int64_t*)(dst)) = *((int64_t*)(src));
		break;
	case Btype::BOOL:
		*((int8_t*)(dst)) = *((int8_t*)(src));
		break;
	default:
		return 1;
		break;
	}
	return 0;
}
static inline int elecpy(void *dst, const void *src, HIRBtype ty)
{
	switch (ty)
	{
	case HIRBtype::FLOAT: case HIRBtype::I32:
		*((int32_t*)(dst)) = *((int32_t*)(src));
		break;
	case HIRBtype::DOUBLE:
		*((int64_t*)(dst)) = *((int64_t*)(src));
		break;
	case HIRBtype::I8:
		*((int8_t*)(dst)) = *((int8_t*)(src));
		break;
	default:
		return 1;
		break;
	}
	return 0;
}

void * get_data_of_const_init_val(CACTParser::Const_init_valContext *ctx,
	const std::vector<uint32_t> &dim, bool is_external = false,
	void *data = nullptr, size_t offset = 0u, uint32_t level = 0u);

//HIRTempno get_dest_of_inst(const HIRInst &inst);
//size_t get_srctemps_of_inst(const HIRInst &inst, std::vector<HIRTempno> &srctemps);

std::ostream & operator<<(std::ostream & os, const extvar_t &xvar);
std::ostream & operator<<(std::ostream & os, const HIRFunct &funct);

#endif